#!/usr/bin/env python

import numpy as np
from sklearn.base import BaseEstimator, ClassifierMixin
from sklearn.cluster import KMeans, kmeans_plusplus
from sklearn.metrics.pairwise import euclidean_distances
from sklearn.model_selection import train_test_split
from sklearn import datasets
from semi_kmeans import *


digists = datasets.load_digits()
X_train, X_test, y_train, y_test = train_test_split(digists.data, digists.target, test_size=0.5)

X_labeled, X_unlabeled, y_labeled, _ = train_test_split(X_train, y_train, test_size=0.95)

if __name__ == '__main__':
    
    km = SemiKMeans(n_clusters=10)
    km.fit(X_labeled, y_labeled, X_unlabeled) # y_test0 is unknown
    skm = SupervisedKMeans()
    skm.fit(X_labeled, y_labeled)
    print(f"""
    # clusters: 10
    # samples: {X_labeled.shape[0]} + {X_unlabeled.shape[0]}

    SemiKMeans: {km.score(X_test, y_test)}
    SupervisedKMeans: {skm.score(X_test, y_test)}
    """)


